package org.zoomdev.zoom.asm;

import com.alibaba.fastjson2.JSONWriter;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Label;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.GeneratorAdapter;
import org.objectweb.asm.commons.Method;
import org.zoomdev.zoom.common.exceptions.ZoomException;
import org.zoomdev.zoom.security.Sensitive;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MethodBuilder {
    // 关联类信息
    private final ClassBuilder classBuilder;

    // 方法基本信息
    private final String methodName;
    // 指令列表
    private final List<Instruction> instructions = new ArrayList<>();
    // 标签管理
    private final Map<String, Label> labels = new HashMap<>();
    private Type returnType = Type.VOID_TYPE;
    private Type[] paramTypes = new Type[0];
    private int access = Opcodes.ACC_PUBLIC;
    // 局部变量管理
    private int nextLocalVarIndex = 0;
    // 是否已经构建标志
    private boolean built = false;


    public MethodBuilder(ClassBuilder classBuilder, String methodName) {
        this.classBuilder = classBuilder;
        this.methodName = methodName;
    }

    // === 方法基本信息设置 ===

    public MethodBuilder withAccess(int access) {
        this.access = access;
        return this;
    }

    public MethodBuilder returns(Class<?> returnType) {
        this.returnType = Type.getType(returnType);
        return this;
    }

    public MethodBuilder returns(Type returnType) {
        this.returnType = returnType;
        return this;
    }

    public MethodBuilder params(Class<?>... paramClasses) {
        this.paramTypes = new Type[paramClasses.length];
        for (int i = 0; i < paramClasses.length; i++) {
            this.paramTypes[i] = Type.getType(paramClasses[i]);
        }
        nextLocalVarIndex = paramClasses.length + (isStatic() ? 0 : 1); // 保留this引用
        return this;
    }

    private String[] paramNames;
    public MethodBuilder paramsAndNames(Object ... objects) {
        if(objects.length%2!=0){
            throw new ZoomException("必须是2的倍数");
        }
        this.paramTypes = new Type[objects.length/2];
        this.paramNames = new String[objects.length/2];
        for (int i = 0; i < objects.length; i+=2) {
            this.paramTypes[i] = Type.getType((Class)objects[i]);
            this.paramNames[i+1] = (String) objects[i+1];
        }
        nextLocalVarIndex = objects.length/2 + (isStatic() ? 0 : 1); // 保留this引用
        return this;
    }

    public MethodBuilder params(Type... paramTypes) {
        this.paramTypes = paramTypes;
        nextLocalVarIndex = paramTypes.length + (isStatic() ? 0 : 1); // 保留this引用
        return this;
    }

    // === 方法指令添加 ===

    public MethodBuilder addInstruction(Instruction instruction) {
        instructions.add(instruction);
        return this;
    }

    // 添加一组指令序列
    public MethodBuilder addInstructions(Instruction... instructions) {
        for (Instruction instruction : instructions) {
            this.instructions.add(instruction);
        }
        return this;
    }

    // === 便捷指令方法 ===

    // 加载常量值
    public MethodBuilder loadConstant(Object value) {
        return addInstruction(Instructions.push(value));
    }

    // 加载this引用
    public MethodBuilder loadThis() {
        if (isStatic()) {
            throw new IllegalStateException("Cannot load 'this' in a static method");
        }
        return addInstruction(Instructions.loadThis());
    }



    // 加载方法参数
    public MethodBuilder loadArg(int index) {
        return addInstruction(Instructions.loadArg(index));
    }

    // 加载局部变量
    public MethodBuilder loadLocal(int index) {
        return addInstruction(Instructions.loadLocal(index));
    }

    public MethodBuilder getField(String fieldName, Type fieldType) {
        return addInstruction(Instructions.getField(
                classBuilder.getClassType(),
                fieldName,
                fieldType
        ));
    }
    // 获取字段
    public MethodBuilder getField(String fieldName, Class<?> fieldType) {
        return addInstruction(Instructions.getField(
                classBuilder.getClassName(),
                fieldName,
                fieldType
        ));
    }

    public MethodBuilder swap() {
        return addInstruction(mg -> mg.visitInsn(Opcodes.SWAP));
    }

    public MethodBuilder putField(String fieldName, Type fieldType) {
        return addInstruction(Instructions.putField(
                classBuilder.getClassType(),
                fieldName,
                fieldType
        ));
    }

    // 设置字段
    public MethodBuilder putField(String fieldName, Class<?> fieldType) {
        return addInstruction(Instructions.putField(
                classBuilder.getClassName(),
                fieldName,
                fieldType
        ));
    }


    // 调用父类方法
    public MethodBuilder invokeSuper(
            String methodName,
            Class<?> returnType,
            Class<?>... paramTypes) {
        return addInstruction(Instructions.invokeVirtual(
                classBuilder.getSuperClass(),
                methodName,
                returnType,
                paramTypes
        ));
    }

    // === 控制流程 ===

    // 添加标签
    public MethodBuilder defineLabel(String name) {
        Label label = new Label();
        labels.put(name, label);
        return this;
    }

    // 跳转到标签
    public MethodBuilder goTo(String labelName) {
        Label label = getLabel(labelName);
        return addInstruction(Instructions.goTo(label));
    }

    // === 局部变量管理 ===

    // 分配新的局部变量
    public int allocateLocalVar(Class<?> type) {
        int size = Type.getType(type).getSize();
        int index = nextLocalVarIndex;
        nextLocalVarIndex += size;
        return index;
    }

    // === 构建方法 ===
    public String getMethodName() {
        return methodName;
    }

    public void build(ClassWriter cw) {
        if (built) {
            throw new IllegalStateException("Method already built: " + methodName);
        }
        built = true;
        // 初始化局部变量索引
        boolean isStatic = (access & Opcodes.ACC_STATIC) != 0;
        initLocalVars(isStatic, paramTypes);

        Method method = new Method(methodName, returnType, paramTypes);
        GeneratorAdapter mg = new GeneratorAdapter(access, method, null, null, cw);
        mg.visitCode();

        // 应用所有指令
        for (Instruction instruction : instructions) {
            instruction.apply(mg);
        }


        // ... 创建 GeneratorAdapter 等 ...
        for (LocalVarInfo var : localVars.values()) {
            if (var.startLabel != null && var.endLabel != null) {
                mg.visitLocalVariable(
                        var.name,
                        var.type.getDescriptor(),
                        null, // 泛型签名（可选）
                        var.startLabel,
                        var.endLabel,
                        var.index
                );
            }
        }


        mg.endMethod();
    }
    // === 辅助方法 ===

    public MethodBuilder returnValue() {
        return addInstruction(new ReturnInstruction());
    }

    public MethodBuilder push(Object value) {
        return addInstruction(mg -> {
            if (value instanceof Integer) mg.push((Integer) value);
            else if (value instanceof Long) mg.push((Long) value);
            else if (value instanceof Float) mg.push((Float) value);
            else if (value instanceof Double) mg.push((Double) value);
            else if (value instanceof String) mg.push((String) value);
            else if (value instanceof Boolean) mg.push((Boolean) value);
            else if (value instanceof Character) mg.push((Character) value);
            else if (value instanceof Type) mg.push((Type) value);
            else if (value == null) mg.visitInsn(Opcodes.ACONST_NULL);
            else throw new IllegalArgumentException("Unsupported type: " + value.getClass());
        });
    }

    public MethodBuilder getField(String ownerClass, String fieldName, Class<?> fieldType) {
        String internalName = normalizeClassName(ownerClass);
        String descriptor = Type.getType(fieldType).getDescriptor();

        return addInstruction(mg -> mg.getField(
                Type.getObjectType(internalName),
                fieldName,
                Type.getType(descriptor)));
    }

    public MethodBuilder putField(String ownerClass, String fieldName, Class<?> fieldType) {
        String internalName = normalizeClassName(ownerClass);
        String descriptor = Type.getType(fieldType).getDescriptor();

        return addInstruction(mg -> mg.putField(
                Type.getObjectType(internalName),
                fieldName,
                Type.getType(descriptor)));
    }

    public MethodBuilder getStatic(String ownerClass, String fieldName, Class<?> fieldType) {
        String internalName = normalizeClassName(ownerClass);
        String descriptor = Type.getType(fieldType).getDescriptor();

        return addInstruction(mg -> mg.getStatic(
                Type.getObjectType(internalName),
                fieldName,
                Type.getType(descriptor)));
    }

    public MethodBuilder putStatic(String ownerClass, String fieldName, Class<?> fieldType) {
        String internalName = normalizeClassName(ownerClass);
        String descriptor = Type.getType(fieldType).getDescriptor();

        return addInstruction(mg -> mg.putStatic(
                Type.getObjectType(internalName),
                fieldName,
                Type.getType(descriptor)));
    }

    // === 字段访问 ===
    public MethodBuilder invokeVirtual(
            Type owner,
            String methodName,
            Class<?> returnType,
            Class<?>... paramTypes) {

        return invokeVirtual(
                owner,
                methodName,
                Type.getType(returnType),
                getTypes(paramTypes));
    }
    public MethodBuilder invokeVirtual(
            Class<?> owner,
            String methodName,
            Class<?> returnType,
            Class<?>... paramTypes) {

        return invokeVirtual(
                Type.getType(owner),
                methodName,
                Type.getType(returnType),
                getTypes(paramTypes));
    }

    public MethodBuilder invokeVirtual(
            Type ownerType,
            String methodName,
            Type returnType,
            Type... paramTypes) {

        return addInstruction(mg -> mg.invokeVirtual(
                ownerType,
                new Method(methodName, returnType, paramTypes)));
    }

    public MethodBuilder invokeInterface(
            Class<?> owner,
            String methodName,
            Class<?> returnType,
            Class<?>... paramTypes) {

        return invokeInterface(
                Type.getType(owner),
                methodName,
                Type.getType(returnType),
                getTypes(paramTypes));
    }

    public MethodBuilder invokeInterface(
            Type ownerType,
            String methodName,
            Type returnType,
            Type... paramTypes) {

        return addInstruction(mg -> mg.invokeInterface(
                ownerType,
                new Method(methodName, returnType, paramTypes)));
    }

    // === 方法调用 ===

    public MethodBuilder invokeStatic(
            Class<?> owner,
            String methodName,
            Class<?> returnType,
            Class<?>... paramTypes) {

        return invokeStatic(
                Type.getType(owner),
                methodName,
                Type.getType(returnType),
                getTypes(paramTypes));
    }

    public MethodBuilder invokeStatic(
            Type ownerType,
            String methodName,
            Type returnType,
            Type... paramTypes) {

        return addInstruction(mg -> mg.invokeStatic(
                ownerType,
                new Method(methodName, returnType, paramTypes)));
    }

    public MethodBuilder invokeConstructor(Class<?>... paramTypes) {
        return loadThis().invokeConstructor(
                classBuilder.getSuperClassType(),
                getTypes(paramTypes));
    }

    public MethodBuilder invokeConstructor(Type ownerType, Type... paramTypes) {
        return addInstruction(mg -> mg.invokeConstructor(
                ownerType,
                new Method("<init>", Type.VOID_TYPE, paramTypes)));
    }

    public MethodBuilder newInstance(Class<?> type) {
        return newInstance(Type.getType(type));
    }

    public MethodBuilder newInstance(Type type) {
        return addInstruction(mg -> mg.newInstance(type));
    }

    public MethodBuilder newArray(Class<?> componentType) {
        return newArray(Type.getType(componentType));
    }

    public MethodBuilder newArray(Type componentType) {
        return addInstruction(mg -> mg.newArray(componentType));
    }

    // === 对象创建 ===
//    public MethodBuilder arrayLoad(Type componentType) {
//
//        if (componentType == byte.class || componentType == boolean.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.BALOAD));
//        } else if (componentType == char.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.CALOAD));
//        } else if (componentType == short.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.SALOAD));
//        } else if (componentType == int.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.IALOAD));
//        } else if (componentType == float.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.FALOAD));
//        } else if (componentType == long.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.LALOAD));
//        } else if (componentType == double.class) {
//            return addInstruction(mg -> mg.visitInsn(Opcodes.DALOAD));
//        } else { // 对象类型
//            return addInstruction(mg -> mg.visitInsn(Opcodes.AALOAD));
//        }
//    }
    public MethodBuilder arrayLoad(Class<?> componentType) {
        if (componentType == byte.class || componentType == boolean.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.BALOAD));
        } else if (componentType == char.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.CALOAD));
        } else if (componentType == short.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.SALOAD));
        } else if (componentType == int.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.IALOAD));
        } else if (componentType == float.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.FALOAD));
        } else if (componentType == long.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.LALOAD));
        } else if (componentType == double.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.DALOAD));
        } else { // 对象类型
            return addInstruction(mg -> mg.visitInsn(Opcodes.AALOAD));
        }
    }

    public MethodBuilder arrayStore(Class<?> componentType) {
        if (componentType == byte.class || componentType == boolean.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.BASTORE));
        } else if (componentType == char.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.CASTORE));
        } else if (componentType == short.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.SASTORE));
        } else if (componentType == int.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.IASTORE));
        } else if (componentType == float.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.FASTORE));
        } else if (componentType == long.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.LASTORE));
        } else if (componentType == double.class) {
            return addInstruction(mg -> mg.visitInsn(Opcodes.DASTORE));
        } else { // 对象类型
            return addInstruction(mg -> mg.visitInsn(Opcodes.AASTORE));
        }
    }

    // === 数组操作 ===

    public MethodBuilder math(int opcode, Class<?> type) {
        return addInstruction(mg -> mg.math(opcode, Type.getType(type)));
    }

    public MethodBuilder goTo(Label label) {
        return addInstruction(mg -> mg.goTo(label));
    }

    public MethodBuilder ifNull(Label label) {
        return addInstruction(mg -> mg.ifNull(label));
    }

    public MethodBuilder ifNonNull(Label label) {
        return addInstruction(mg -> mg.ifNonNull(label));
    }

    // === 数学运算 ===

    public MethodBuilder ifCmp(Type type, int mode, Label label) {
        return addInstruction(mg -> mg.ifCmp(type, mode, label));
    }



    ////////////////////////////////////////////////

    // 局部变量管理
    private final Map<String, LocalVarInfo> localVars = new HashMap<>();


    // 存储局部变量信息
    private static class LocalVarInfo {
        final int index;
        final Type type;
        final String name;
        Label startLabel;
        Label endLabel;
        LocalVarInfo(int index, Type type, String name) {
            this.index = index;
            this.type = type;
            this.name = name;
        }
    }

    // === 局部变量管理 ===

    /**
     * 存储值到局部变量
     * @param name 变量名
     * @param type 变量类型
     * @return 当前 MethodBuilder 实例
     */
    public MethodBuilder storeLocal(String name, Class<?> type) {
        return storeLocal(name, Type.getType(type));
    }

    public MethodBuilder storeLocal(String name, Type type) {
        // 分配局部变量索引
        int index0 = allocateLocal(type);
        localVars.put(name, new LocalVarInfo(index0, type, name));

        // 生成存储指令
        return addInstruction(mg -> {
            int index = index0;
            switch (type.getSort()) {
                case Type.BOOLEAN:
                case Type.BYTE:
                case Type.CHAR:
                case Type.SHORT:
                case Type.INT:
                    mg.storeLocal(index, type);
                    break;
                case Type.LONG:
                    mg.storeLocal(index, Type.LONG_TYPE);
                    break;
                case Type.FLOAT:
                    mg.storeLocal(index, Type.FLOAT_TYPE);
                    break;
                case Type.DOUBLE:
                    mg.storeLocal(index, Type.DOUBLE_TYPE);
                    break;
                case Type.ARRAY:
                case Type.OBJECT:
                    mg.storeLocal(index, type);
                    break;
                default:
                    throw new IllegalArgumentException("Unsupported type: " + type);
            }
        });
    }

    /**
     * 从局部变量加载值
     * @param name 变量名
     * @return 当前 MethodBuilder 实例
     */
    public MethodBuilder loadLocal(String name) {
        LocalVarInfo var0 = localVars.get(name);
        if (var0 == null) {
            throw new IllegalArgumentException("Local variable not defined: " + name);
        }

        return addInstruction(mg -> {
            LocalVarInfo var = var0;
            switch (var.type.getSort()) {
                case Type.BOOLEAN:
                case Type.BYTE:
                case Type.CHAR:
                case Type.SHORT:
                case Type.INT:
                    mg.loadLocal(var.index, var.type);
                    break;
                case Type.LONG:
                    mg.loadLocal(var.index, Type.LONG_TYPE);
                    break;
                case Type.FLOAT:
                    mg.loadLocal(var.index, Type.FLOAT_TYPE);
                    break;
                case Type.DOUBLE:
                    mg.loadLocal(var.index, Type.DOUBLE_TYPE);
                    break;
                case Type.ARRAY:
                case Type.OBJECT:
                    mg.loadLocal(var.index, var.type);
                    break;
                default:
                    throw new IllegalArgumentException("Unsupported type: " + var.type);
            }
        });
    }

    /**
     * 分配局部变量索引
     * @param type 变量类型
     * @return 分配的索引
     */
    private int allocateLocal(Type type) {
        int size = type.getSize();
        int index = nextLocalVarIndex;
        nextLocalVarIndex += size;
        return index;
    }

    /**
     * 初始化局部变量索引（在方法构建开始时调用）
     * @param isStatic 是否是静态方法
     * @param paramTypes 参数类型列表
     */
    private void initLocalVars(boolean isStatic, Type[] paramTypes) {
        // 重置索引
//        nextLocalVarIndex = 0;
//        localVars.clear();
//
//        // 非静态方法：索引0是this
//        if (!isStatic) {
//            nextLocalVarIndex += Type.getType(Object.class).getSize(); // this引用
//        }
//
//        // 处理参数
//        for (Type paramType : paramTypes) {
//            nextLocalVarIndex += paramType.getSize();
//        }
    }
  /**
          * 设置局部变量作用域
     */
    public MethodBuilder setLocalVarScope(String name, String startLabel, String endLabel) {
        LocalVarInfo var = localVars.get(name);
        if (var == null) {
            throw new IllegalArgumentException("Local variable not defined: " + name);
        }

        var.startLabel = getLabel(startLabel);
        var.endLabel = getLabel(endLabel);
        return this;
    }


    public MethodBuilder storeLocal(int index){
        return addInstruction(mg -> mg.storeLocal(index));
    }
    // === 控制流 ===

    /**
     *
     * @return
     */
    private final Map<String, Type> typeCache = new HashMap<>();

    public MethodBuilder checkcast(Class<?> type) {
        return checkcast(Type.getType(type));
    }

    public MethodBuilder checkcast(Type type) {
        return addInstruction(mg -> {
            if (type.getSort() == Type.OBJECT || type.getSort() == Type.ARRAY) {
                mg.checkCast(type);
            } else {
                // 对于基本类型，不需要 checkcast
                // 但可以添加类型转换指令（如 i2l, f2d 等）
                // 这里只处理对象类型
            }
        });
    }

    public MethodBuilder checkcast(String className) {
        // 处理不同格式的类名
        String internalName = className.replace('.', '/');
        if (!internalName.startsWith("L") && !internalName.endsWith(";")) {
            internalName = "L" + internalName + ";";
        }
        return checkcast(Type.getObjectType(internalName));
    }

//    public MethodBuilder checkcastGeneric(Class<?> rawType, Class<?>... typeArguments) {
//        Type[] argTypes = new Type[typeArguments.length];
//        for (int i = 0; i < typeArguments.length; i++) {
//            argTypes[i] = Type.getType(typeArguments[i]);
//        }
//        return checkcastGeneric(Type.getType(rawType), argTypes);
//    }

//    public MethodBuilder checkcastGeneric(Type rawType, Type... typeArguments) {
//        StringBuilder signature = new StringBuilder();
//        signature.append("L").append(rawType.getInternalName());
//
//        if (typeArguments.length > 0) {
//            signature.append("<");
//            for (Type arg : typeArguments) {
//                signature.append(arg.getDescriptor());
//            }
//            signature.append(">");
//        }
//
//        signature.append(";");
//
//        return addInstruction(mg -> {
//            mg.visitTypeInsn(Opcodes.CHECKCAST, rawType.getInternalName());
//            // 添加泛型签名信息
//            mg.visitFrame(Opcodes.F_NEW, 0, null, 0, null);
//            mg.visitLocalVariable("this", signature.toString(), null, null, null);
//        });
//    }




    private boolean isStatic() {
        return (access & Opcodes.ACC_STATIC) != 0;
    }

    private Label getLabel(String name) {
        Label label = labels.get(name);
        if (label == null) {
            throw new IllegalArgumentException("Label '" + name + "' not defined");
        }
        return label;
    }

    // === 类型转换 ===

    private boolean isReturnStatementPresent() {
        for (Instruction instruction : instructions) {
            System.out.println(instruction.getClass());
            if (instruction instanceof ReturnInstruction) {
                return true;
            }
        }
        return false;
    }

    private Type[] getTypes(Class<?>[] classes) {
        Type[] types = new Type[classes.length];
        for (int i = 0; i < classes.length; i++) {
            types[i] = Type.getType(classes[i]);
        }
        return types;
    }

    // === 实用方法 ===

    private String normalizeClassName(String className) {
        if (className.contains(".")) {
            return className.replace('.', '/');
        }
        return className;
    }

    ////////////////////////////////////////////////////

    /**
     * 高级封装， 表示  this.fieldName = param.index
     * @param fieldName
     * @param index
     * @return
     */
    public MethodBuilder assign(String fieldName,int index){
        return loadArg(index)
                .putField(fieldName, classBuilder.getField(fieldName).getFieldType());
    }

    public MethodBuilder putField(String fieldName) {
        return addInstruction(Instructions.putField(
                classBuilder.getClassType(),
                fieldName,
                classBuilder.getField(fieldName).getFieldType()
        ));
    }
    private String ref;

    public MethodBuilder refThis(){
        return loadThis();
    }


    public MethodBuilder call(int index,String method,Class returnType,Class...types){
        return loadArg(index) // JSONWriter
                .invokeVirtual(paramTypes[index], method, returnType,types);
    }

//    public MethodBuilder loadThisArray(String name,int index){
//        return loadThis()
//                .getField(name, classBuilder.getField(name).getFieldType())
//                .push(0)
//                .arrayLoad(classBuilder.getField(name).getType());         //Sensitive se = this.sensitives[0]
//    }



//
//
//    public MethodBuilder loadThis() {
//        return addInstruction(mg -> mg.loadThis());
//    }
//
//    public MethodBuilder loadArg(int index) {
//        return addInstruction(mg -> mg.loadArg(index));
//    }
//
//
//    // 调用实例方法
//    public MethodBuilder invokeVirtual(
//            Class<?> owner,
//            String methodName,
//            Class<?> returnType,
//            Class<?>... paramTypes) {
//        return addInstruction(Instructions.invokeVirtual(owner, methodName, returnType, paramTypes));
//    }
//
//    // 调用静态方法
//    public MethodBuilder invokeStatic(
//            Class<?> owner,
//            String methodName,
//            Class<?> returnType,
//            Class<?>... paramTypes) {
//        return addInstruction(Instructions.invokeStatic(owner, methodName, returnType, paramTypes));
//    }
}